基于加速平台构建的强大优势: DeepVariant 如何使用 Intel 的 AVX-512 优化程序
Google Brain Genomics 团队开发的 DeepVariant 是一种用于分析个体 DNA 序列的开源工具,基于 TensorFlow 构建(https://github.com/google/deepvariant)。
相较于上一版 DeepVariant (v0.6),最新版本 (v0.7) 不仅将端到端速度提升三倍,还能相应降低成本。这种速度提升很大程度上得益于 DeepVariant 能够利用全新 Intel® 高级矢量扩展 (AVX-512) 指令集。对 AVX-512 的使用诠释了系统和平台方面的最新进展如何帮助加速并影响实用研究。
什么是 AVX-512 指令集和 TensorFlow 数学内核库 (TF-MKL)?
Intel CPU 广泛应用于 Google Cloud Platform,为许多大型独立系统提供计算执行。处理单元包含寄存器,系统可以将数据存放其中。处理单元每次只能对寄存器中的内容执行一种类型的操作。因此,在组织需要相同操作的项并将其加载至寄存器中时,我们需同时对这两个过程进行大量优化。
您可以将此想象成活跃的交通枢纽,公共汽车不断从这里发出,前往不同的目的地。此枢纽必须确定派出送客的公共汽车,尽量使每辆车都能满座,避免乘客等待太久而延误行程。
寄存器的大小及其有意义地存放数据的能力会限制处理的速度。AVX-512 指令集是一种经过专门设计的方法,其会封装此寄存器,使处理器能够一次处理更多信息。在此类比中,该指令集相当于派遣容量较大的公共汽车前往热门目的地。如需了解更多详情,请参阅 Intel 文档(https://www.intel.com/content/www/us/en/architecture-and-technology/avx-512-overview.html)。
在 TensorFlow 中训练和应用深度学习模型的过程涉及到大量相同类型的矢量/矩阵运算。借助 Intel 开发的 适用于深度神经网络的 Intel 数学内核库,TensorFlow 应用可以高效填充这些容量较大的寄存器。相应地,Intel CPU 也能为这些应用提供速度更快的计算。
图 1. AVX-512 在 TF-MKL 中的优势(简化概念示意图)
量化 TF-MKL 和 Intel AVX-512 对 DeepVariant 的影响
在调用测序数据变体的过程中,DeepVariant 会执行三个不同的步骤。如需更多详情,请参阅 此刊发论文 或 Google AI 博文。计算最密集的步骤是 call_variants 阶段,该阶段会利用卷积神经网络来界定个体基因数据与参考基因数据中的位置是否存在差异。这些结果既可在临床领域用于诊断和治疗患者,也可在研究中用于探索和开发新药。
注:此刊发论文 链接
https://www.nature.com/articles/nbt.4235Google AI 博文 链接https://ai.googleblog.com/2017/12/deepvariant-highly-accurate-genomes.html
如果是旧版本的 DeepVariant (v0.6),则在配备 Intel Skylake CPU 的 64 核设备上完成 call_variants 步骤需耗时 11 小时 11 分钟。至于最新版的 DeepVariant (v0.7),我们优化了其使用 TensorFlow 的过程,以充分利用新的 TF-MKL 库。利用这些创新,在相同的 64 核设备上执行 call_variants 的运行时间会减少为 3 小时 25 分钟。
鉴于 DeepVariant v0.6 已经使用了 Intel 核心库的一些优化程序,因此为公平量化 TF-MKL 带来的加速效果,我们决定使用不含优化程序的 TensorFlow 库来重新构建版本完全相同的 DeepVariant (v0.7)。基于这一配置,在相同的 64 核设备上执行 call_variants 需耗时 14 小时 7 分钟。
图 2. 通过使用 TF-MKL 在 DeepVariant v0.7 中实现运行时加速
您可以在 DeepVariant GitHub 页面找到相关说明,了解如何使用预装所有适用库的 Docker 容器,以便在您的本地或云环境中运行 DeepVariant,并充分利用这一速度提升的优势。
注:DeepVariant GitHub 链接
https://ai.googleblog.com/2017/12/deepvariant-highly-accurate-genomes.html
这种速度提升及其他优化(我们希望日后另作详述)都会相应降低成本。在针对典型覆盖全基因组的基准测试中,DeepVariant v0.7 在 Google Cloud Platform 上的运行成本为 2 到 3 美元,相较于 DeepVariant v0.6 减少了 3 倍。
基于快速改进的平台构建的强大优势
本文详细介绍了我们如何继续提升 DeepVariant 的运行速度(以及如何提升其精确度和扩展性)。由于开发者和用户群的不懈努力,我们得以使用这些自带持续改进优势的技术减轻工作负担。正是得益于 TensorFlow 的广泛采用,Google 内外部的许多团队才能积极投入工作,从而提升 DeepVariant 的速度和性能。Intel 为 AVX-512 加速 TensorFlow 所做的工作便是这一成效的精彩范例。
此外,随着新一代 Intel CPU 的纷纷涌现,AVX-512 的性能获得了显著提升。在运行 DeepVariant 时,新型 Intel Skylake CPU 的速度胜于 Intel Broadwell CPU。我们预期,随着向量运算的并行化变得愈加重要,这种趋势会愈发明显。因此,我们也期待 DeepVariant 的速度在未来几代 Intel CPU 中获得自然提升,甚至以超越一般应用的速度提升。
这就是基于快速改进(无论是硬件、软件,抑或是介于二者之间的技术)的平台构建的优势。这就是为何站在基因组学与信息学的交汇之处会让人如此兴奋。随着对基因组进行排序的需求呈指数增长,我们必须对分析这些基因组的软硬件进行升级,以满足这种需求。
更多 AI 相关阅读: